# -*- coding: UTF-8 -*-

"""classify.py"""
import sys
import os
import acl



from constants import ACL_MEM_MALLOC_HUGE_FIRST, ACL_MEMCPY_DEVICE_TO_DEVICE, IMG_EXT
from acllite_model import AclLiteModel
from acllite_image import AclLiteImage
from acllite_resource import AclLiteResource
from resnet50_classes_age import get_resnet50_class
from PIL import Image, ImageDraw, ImageFont
import numpy as np


MODEL_WIDTH = 224
MODEL_HEIGHT = 224


MODEL_PATH = '../model/resnet50_age.om'
MODEL_WIDTH = 224
MODEL_HEIGHT = 224


class Classify(object):
    """classify"""

    def __init__(self, model_path, model_width, model_height):
        self._model_path = model_path
        self._model_width = model_width
        self._model_height = model_height
        self._model = AclLiteModel(model_path)

    def __del__(self):
        print("[Sample] class Samle release source success")

    def pre_process(self, image):
        """preprocess"""
        input_image = Image.open(image)
        input_image = input_image.resize((448, 448))
        # hwc
        img = np.array(input_image)
        height = img.shape[0]
        width = img.shape[1]
        h_off = int((height - 224) / 2)
        w_off = int((width - 224) / 2)
        crop_img = img[h_off:height - h_off, w_off:width - w_off, :]
        # rgb to bgr
        print("crop shape = ", crop_img.shape)
        img = crop_img[:, :, ::-1]
        shape = img.shape
        print("img shape = ", shape)
        img = img.astype("float32")
        img[:, :, 0] *= 0.003922
        img[:, :, 1] *= 0.003922
        img[:, :, 2] *= 0.003922
        img[:, :, 0] -= 0.4914
        img[:, :, 0] = img[:, :, 0] / 0.2023
        img[:, :, 1] -= 0.4822
        img[:, :, 1] = img[:, :, 1] / 0.1994
        img[:, :, 2] -= 0.4465
        img[:, :, 2] = img[:, :, 2] / 0.2010
        img = img.reshape([1] + list(shape))
        # nhwc -> nchw
        result = img.transpose([0, 3, 1, 2]).copy()
        return result

    def inference(self, resized_image):
        """inference"""
        return self._model.execute([resized_image, ])

    def post_process(self, infer_output, image_file):
        """postprocess"""
        print("post process")
        data = infer_output[0]
        print("data shape = ", data.shape)
        vals = data.flatten()
        max = 0
        sum = 0
        for i in range(0, 10):
            if vals[i] > max:
                max = vals[i]
        for i in range(0, 10):
            vals[i] = np.exp(vals[i] - max)
            sum += vals[i]
        for i in range(0, 10):
            vals[i] /= sum
        print("vals shape = ", vals.shape)
        top_k = vals.argsort()[-1:-6:-1]
        print("images:{}".format(image_file))
        print("======== top5 inference results: =============")
        for n in top_k:
            object_class = get_resnet50_class(n)
            print("label:%d  confidence: %f, class: %s" % (n, vals[n], object_class))

        # using pillow, the category with the highest confidence is written on the image and saved locally
        if len(top_k):
            object_class = get_resnet50_class(top_k[0])
            output_path = '../img/resnet50_age_rqy/resnet50_age.jpg'
            origin_img = Image.open(image_file)
            draw = ImageDraw.Draw(origin_img)
            font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", size=20)
            draw.text((10, 50), object_class, font=font, fill=255)
            origin_img.save(output_path)


def resnet50(MODEL_PATH,img_path):
    
    
    # Instantiation classification detection, incoming om model path, model input width and height parameters
    classify = Classify(MODEL_PATH, MODEL_WIDTH, MODEL_HEIGHT)

    

    
    # preprocess image
    resized_image = classify.pre_process(img_path)
    print("pre process end")
    # inference
    result = classify.inference(resized_image)
    # post process
    classify.post_process(result, img_path)



